#
# SSH Terminal implementation
#

from twisted.cred import portal, checkers, credentials
from twisted.conch import error, avatar, recvline, interfaces as conchinterfaces
from twisted.conch.ssh import factory, userauth, connection, keys, session, common
from twisted.conch.insults import insults
from twisted.application import service, internet
from zope.interface import implements
import logging

from fgs.proxy.error import FgsProxyError

class FgsSSHProtocol(recvline.HistoricRecvLine):
    def __init__(self, user, proxy):
        self._user = user
        self._proxy = proxy
        self._log = logging.getLogger('ssh.terminal')
        self._motd = proxy.config.get('ssh', 'motd_file')

        from os.path import exists
        if not exists(self._motd):
            self._log.error("MOTD file %s not found", self._motd)
        else:
            name = proxy.config.get('proxy', 'name')
            f = file(self._motd)
            self._motd = f.read() % {'name': name}
            f.close()

    def _construct_commandlist(self):
        public_methods = filter(
            lambda funcname: funcname.startswith('do_'), dir(self))
        commands = [cmd[3:] for cmd in public_methods]
        commands.sort()

        missing_len = len(commands) % 4
        if missing_len != 0:
            commands += ["" for _ in range(4 - (len(commands) % 4))]
        commands = [commands[x:x+len(commands)/4]
                    for x in xrange(0, len(commands), len(commands)/4)]
        print commands
        commands = zip(*commands)
        print commands
        self._command_list = ["  %-16s%-16s%-16s%-16s" % l
                               for l in commands]

    def connectionMade(self):
        self._log.info("%s logged in" % self._user.username)
        recvline.HistoricRecvLine.connectionMade(self)
        self.keyHandlers.update({
            '\x04': self.do_quit
            })
        for line in self._motd.split('\n'):
            self.terminal.write(line)
            self.terminal.nextLine()
        self._construct_commandlist()
        self.do_help()
        self.showPrompt()

    def connectionLost(self, reason):
        self._log.info("%s logged out" % self._user.username)

    def showPrompt(self):
        self.terminal.write("% ")

    def getCommandFunc(self, cmd):
        return getattr(self, 'do_' + cmd, None)

    def lineReceived(self, line):
        line = line.strip()
        if line:
            cmdAndArgs = line.split()
            cmd = cmdAndArgs[0]
            args = cmdAndArgs[1:]
            func = self.getCommandFunc(cmd)
            if func:
                try:
                    func(*args)
                except Exception, e:
                    self.terminal.write("Error: %s" % e)
                    self.terminal.nextLine()
            else:
                self.terminal.write("No such command.")
                self.terminal.nextLine()
        self.showPrompt()


    def do_help(self, cmd=''):
        "Get help on a command. Usage: help command"
        if cmd:
            func = self.getCommandFunc(cmd)
            if func:
                self.terminal.write(func.__doc__)
                self.terminal.nextLine()
                return

        self.terminal.write("Available commands:")
        self.terminal.nextLine()
        for l in self._command_list:
            self.terminal.write(l)
            self.terminal.nextLine()
        self.terminal.nextLine()

    def do_quit(self):
        """Ends your session.
        Usage: quit"""
        self.terminal.loseConnection()

    def do_clear(self):
        """Clears the screen.
        Usage: clear"""
        self.terminal.reset()

    def do_shutdown(self):
        """Shutdown the proxy.
        Usage: shutdown"""
        from twisted.internet import reactor

        self._log.info("Shutdown requested by %s", self._user.username)
        self.terminal.write("Shutting down the proxy...")
        self.terminal.nextLine()

        self.do_quit()
        reactor.stop()

class FgsSSHAvatar(avatar.ConchUser):
    implements(conchinterfaces.ISession)

    def __init__(self, username, proxy):
        avatar.ConchUser.__init__(self)
        self._proxy = proxy
        self.username = username
        self.channelLookup.update({'session':session.SSHSession})

    def openShell(self, protocol):
        serverProtocol = insults.ServerProtocol(FgsSSHProtocol, self, self._proxy)
        serverProtocol.makeConnection(protocol)
        protocol.makeConnection(session.wrapProtocol(serverProtocol))

    def getPty(self, terminal, windowSize, attrs):
        return None

    def execCommand(self, protocol, cmd):
        raise NotImplementedError

    def closed(self):
        pass


class FgsSSHRealm(object):
    implements(portal.IRealm)

    def __init__(self, proxy):
        self._proxy = proxy

    def requestAvatar(self, avatarId, mind, *interfaces):
        for iface in interfaces:
            if iface == conchinterfaces.IConchUser:
                return iface, FgsSSHAvatar(avatarId, self._proxy), lambda: None
        raise Exception, "No supported interfaces found."

class FgsSSHPasswordFile(object):
    def __init__(self):
        self._log = logging.getLogger('ssh.cred')

    def _make_hash(self, cleartext, salt=None):
        import sha, base64, random

        if not salt:
            salt = ''.join([chr(random.randint(0,255)) for _ in xrange(6)])
        else:
            salt = base64.decodestring(salt)

        hash_val = base64.encodestring(sha.new(salt+cleartext).digest()).strip()
        hash_val = base64.encodestring(salt).strip() + "$" + hash_val

        return hash_val

    def _read_file(self, filename):
        try:
            f = file(filename)
        except IOError:
            return {}
        userdata = {}
        for line in f:
            fields = line.split(':',1)
            userdata[fields[0]] = fields[1]
        f.close()

        return userdata

    def _write_file(self, filename, userdata):
        f = file(filename, 'w')
        filedata = '\n'.join([':'.join((k,v)) for k,v in userdata.items()])
        f.write(filedata)
        f.write('\n')
        f.close()

    def cred_validate_hash(self, username, passwd, onfile_passwd):
        """Function for the credential validator to validate a
        password against the on-file data"""
        try:
            salt, hash_val = onfile_passwd.split('$')
        except:
            self._log.warn("Password entry for user %s seems corrupted" % username)
            return None

        user_hash = self._make_hash(passwd, salt)
        if user_hash != onfile_passwd:
            self._log.warn("Invalid password for user '%s'" % username)

        return user_hash

    def change_password(self, passwd_file, username, new_pass):
        userdata = self._read_file(passwd_file)
        userdata[username] = self._make_hash(new_pass)
        self._write_file(passwd_file, userdata)

class FgsSSHFactory(factory.SSHFactory):
    def __init__(self, proxy):
        from os.path import exists

        self._log = logging.getLogger('ssh')
        self._log.info("Starting SSH admin console service")

        self.proxy = proxy

        #
        # Check configuration for the SSH subsystem
        #
        self.proxy.config.check_fields(self.proxy.config.check_failure_exception,
                                       ('ssh', 'public_key'),
                                       ('ssh', 'private_key'),
                                       ('ssh', 'passwd_file'),
                                       ('ssh', 'motd_file'))

        passwd_file = self.proxy.config.get('ssh', 'passwd_file')
        pubkeyfile = self.proxy.config.get('ssh', 'public_key')
        privkeyfile = self.proxy.config.get('ssh', 'private_key')

        #
        # Check the password file
        #
        if not exists(passwd_file):
            raise FgsProxyError("The password file '%s'" % passwd_file +
                                " does not exist.")

        #
        # Attempt to load the SSH keys given
        #
        pubkey, privkey = self.getRSAKeys(pubkeyfile, privkeyfile)
        self.publicKeys = {
            'ssh-rsa': keys.getPublicKeyString(data=pubkey) }
        self.privateKeys = {
            'ssh-rsa': keys.getPrivateKeyObject(data=privkey) }

        #
        # Register the portal
        #
        self.portal = portal.Portal(FgsSSHRealm(proxy))
        pass_checker = FgsSSHPasswordFile()
        checker = checkers.FilePasswordDB(passwd_file,
                                          hash=pass_checker.cred_validate_hash)
        self.portal.registerChecker(checker)

    def getRSAKeys(self, pubfile, privfile):
        import os
        if not (os.path.exists(pubfile) and os.path.exists(privfile)):
            # generate a RSA keypair
            from Crypto.PublicKey import RSA
            self._log.warning("SSH server keys missing, regenerating...")

            KEY_LENGTH = 1024
            rsaKey = RSA.generate(KEY_LENGTH, common.entropy.get_bytes)
            publicKeyString = keys.makePublicKeyString(rsaKey)
            privateKeyString = keys.makePrivateKeyString(rsaKey)

            # save keys for next time
            file(pubfile, 'w+b').write(publicKeyString)
            file(privfile, 'w+b').write(privateKeyString)
        else:
            publicKeyString = file(pubfile).read()
            privateKeyString = file(privfile).read()
        self._log.info("Loaded SSH server keys")
        return publicKeyString, privateKeyString
